package core

import (
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/kataras/golog"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"net"
)

func EncodeICMPEchoRequestPacket(id uint16, seq uint16, payload []byte) ([]byte, error) {
	req := &icmp.Message{
		Type: ipv4.ICMPTypeEcho,
		Body: &icmp.Echo{
			ID:   int(id),
			Seq:  int(seq),
			Data: payload,
		},
	}

	data, err := req.Marshal(nil)
	if err != nil {
		golog.Errorf("marshal failed. err: %+v", err)
		return nil, err
	}

	return data, nil
}

func EncodeIPICMPEchoRequestPacket(srcIP string, dstIP string, tos uint8, ttl uint8, id uint16, seq uint16, payload []byte) ([]byte, error) {

	ip := &layers.IPv4{
		SrcIP:    net.ParseIP(srcIP),
		DstIP:    net.ParseIP(dstIP),
		Version:  4,
		Protocol: layers.IPProtocolICMPv4,
		TTL:      ttl,
		TOS:      tos,
	}

	// 构造icmp
	icmpMsg := &layers.ICMPv4{
		Id:       id,
		Seq:      seq,
		TypeCode: layers.ICMPv4TypeEchoRequest<<8 | layers.ICMPv4CodeNet,
	}

	buf := gopacket.NewSerializeBufferExpectedSize(len(payload)+30, 0)
	opts := gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}

	err := gopacket.SerializeLayers(buf, opts, ip, icmpMsg, gopacket.Payload(payload))

	return buf.Bytes(), err
}

func EncodeIPUDPPacket(srcIP string, dstIP string, id uint16, tos uint8, ttl uint8, srcPort uint16, dstPort uint16, payload []byte) ([]byte, error) {

	ip := &layers.IPv4{
		SrcIP:    net.ParseIP(srcIP), // 需要填写实际报文的源ip
		DstIP:    net.ParseIP(dstIP),
		Version:  4,
		Protocol: layers.IPProtocolUDP,
		TTL:      ttl,
		TOS:      tos,
		Id:       id,
	}

	// udp
	udp := &layers.UDP{
		SrcPort: layers.UDPPort(srcPort),
		DstPort: layers.UDPPort(dstPort),
	}
	udp.SetNetworkLayerForChecksum(ip)

	buf := gopacket.NewSerializeBufferExpectedSize(len(payload)+30, 0)
	opts := gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}

	err := gopacket.SerializeLayers(buf, opts, ip, udp, gopacket.Payload(payload))
	return buf.Bytes(), err
}

func EncodeIPTCPSynPacket(srcIP string, dstIP string, id uint16, tos uint8, ttl uint8, srcPort uint16, dstPort uint16, seq uint32, payload []byte) ([]byte, error) {

	ip := &layers.IPv4{
		SrcIP:    net.ParseIP(srcIP),
		DstIP:    net.ParseIP(dstIP),
		Version:  4,
		Protocol: layers.IPProtocolTCP,
		TTL:      ttl,
		TOS:      tos,
		Id:       id,
	}

	/*
		seqID2 := uint32(rand.Intn(math.MaxUint32))
		//log.Debug.Printf("TTL %d, SeqID %d\n", ttl, seqID2)

		sack := layers.TCPOption{
			OptionType:   layers.TCPOptionKindSACKPermitted,
			OptionLength: 2,
		}
		nop := layers.TCPOption{
			OptionType: layers.TCPOptionKindNop,
		}
		itb := make([]byte, 2)
		binary.BigEndian.PutUint16(itb, 1460)
		mss := layers.TCPOption{
			OptionType:   layers.TCPOptionKindMSS,
			OptionLength: 4,
			OptionData:   itb,
		}
		itb = make([]byte, 1)
		itb[0] = 2
		ws := layers.TCPOption{
			OptionType:   layers.TCPOptionKindWindowScale,
			OptionLength: 3,
			OptionData:   itb,
		}
		ts := layers.TCPOption{
			OptionType:   layers.TCPOptionKindTimestamps,
			OptionLength: 10,
			OptionData:   []byte{0xfe, 0xa1, 0xc1, 0xe9, 0x00, 0x00, 0x00, 0x00},
		}

	*/
	tcp := &layers.TCP{
		SrcPort: layers.TCPPort(srcPort), // will be incremented during the scan
		DstPort: layers.TCPPort(dstPort),
		Seq:     seq,
		SYN:     true,
		// CWR:     true,
		// ECE:     true,
		// Seq:     seqID2,
		// Window:  5840,
		// Options: []layers.TCPOption{mss, sack, ts, nop, ws},
	}
	_ = tcp.SetNetworkLayerForChecksum(ip)

	buf := gopacket.NewSerializeBufferExpectedSize(len(payload)+60, 0)
	opts := gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}

	err := gopacket.SerializeLayers(buf, opts, ip, tcp, gopacket.Payload(payload))

	return buf.Bytes(), err
}

func EncodeIPTCPPacket(srcIP string, dstIP string, id uint16, tos uint8, ttl uint8, srcPort uint16, dstPort uint16, seq uint32, payload []byte) ([]byte, error) {

	ip := &layers.IPv4{
		SrcIP:    net.ParseIP(srcIP),
		DstIP:    net.ParseIP(dstIP),
		Version:  4,
		Protocol: layers.IPProtocolTCP,
		TTL:      ttl,
		TOS:      tos,
		Id:       id,
	}

	/*
		seqID2 := uint32(rand.Intn(math.MaxUint32))
		//log.Debug.Printf("TTL %d, SeqID %d\n", ttl, seqID2)

		sack := layers.TCPOption{
			OptionType:   layers.TCPOptionKindSACKPermitted,
			OptionLength: 2,
		}
		nop := layers.TCPOption{
			OptionType: layers.TCPOptionKindNop,
		}
		itb := make([]byte, 2)
		binary.BigEndian.PutUint16(itb, 1460)
		mss := layers.TCPOption{
			OptionType:   layers.TCPOptionKindMSS,
			OptionLength: 4,
			OptionData:   itb,
		}
		itb = make([]byte, 1)
		itb[0] = 2
		ws := layers.TCPOption{
			OptionType:   layers.TCPOptionKindWindowScale,
			OptionLength: 3,
			OptionData:   itb,
		}
		ts := layers.TCPOption{
			OptionType:   layers.TCPOptionKindTimestamps,
			OptionLength: 10,
			OptionData:   []byte{0xfe, 0xa1, 0xc1, 0xe9, 0x00, 0x00, 0x00, 0x00},
		}

	*/
	tcp := &layers.TCP{
		SrcPort: layers.TCPPort(srcPort), // will be incremented during the scan
		DstPort: layers.TCPPort(dstPort),
		Seq:     seq,
		SYN:     false,
		// CWR:     true,
		// ECE:     true,
		// Seq:     seqID2,
		// Window:  5840,
		// Options: []layers.TCPOption{mss, sack, ts, nop, ws},
	}
	_ = tcp.SetNetworkLayerForChecksum(ip)

	buf := gopacket.NewSerializeBufferExpectedSize(len(payload)+60, 0)
	opts := gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}

	err := gopacket.SerializeLayers(buf, opts, ip, tcp, gopacket.Payload(payload))

	return buf.Bytes(), err
}
